/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Contains classes describing angular distribution functions.

// $Id$

#ifndef __angular__
#define __angular__

#include "mcrx-types.h"
#include "ray.h"
#include "constants.h"
#include "boost/shared_ptr.hpp"
#include "blitz/array.h"
#include "mono_poly_abstract.h"

namespace mcrx {
  template <typename> class angular_distribution;
  template<typename> class isotropic_angular_distribution;
  template <typename T> class isotropic_angular_distribution<blitz::Array<T, 1> >; 
  template<typename> class collimated_angular_distribution;
  template <typename T> class collimated_angular_distribution<blitz::Array<T, 1> >; 
  template<typename> class cone_angular_distribution;
  template <typename T> class cone_angular_distribution<blitz::Array<T,1> >;
  template<typename> class ray_distribution_pair;
}


/** Abstract base class for objects that can tell the transfer class
    what the probability is of emission/scattering in certain angles.
    This is necessary for using the next-event estimator.
*/
template <typename T_lambda> 
class mcrx::angular_distribution {
public:
  angular_distribution () {};
  virtual ~angular_distribution () {};
  /** This returns the distribution function. When T_lambda is a blitz
      array, this is less efficient because it requires an
      allocation. However, we can't return an expression because it's
      a virtual function. */
  virtual T_lambda distribution_function(const vec3d& new_direction,
					 const vec3d& old_direction) const= 0;
  /** This function assigns the distribution function in place. This
      is more efficient for arrays since we don't have to allocate the
      return value. */
  virtual void distribution_function(const vec3d& new_direction,
				     const vec3d& old_direction,
				     T_lambda& distr) const= 0;
  virtual boost::shared_ptr<angular_distribution<T_lambda> > clone() const=0;
};


/** Angular_distribution class for isotropic distributions.  Simply
    returns 1/4pi as the distribution function. Used by most of the
    general emission classes. */
template <typename T_lambda> 
class mcrx::isotropic_angular_distribution: 
  public angular_distribution<T_lambda>
{
public:
  isotropic_angular_distribution () {};
  isotropic_angular_distribution (const T_lambda& l): 
    angular_distribution<T_lambda>() {};
  virtual T_lambda distribution_function (const vec3d&, const vec3d&) const {
    return 1/(4*constants::pi); };
  virtual void distribution_function(const vec3d&, const vec3d&,
				     T_lambda& distr) const {
    distr = 1/(4*constants::pi); };
  boost::shared_ptr<angular_distribution<T_lambda> > clone() const {
    return boost::shared_ptr<angular_distribution<T_lambda> >
      (new isotropic_angular_distribution<T_lambda>(*this));};
};

/** Specialization of the isotropic_angular_distribution class for
    polychromatic situations. */
template <typename T> 
class mcrx::isotropic_angular_distribution<blitz::Array<T,1> >: 
  public angular_distribution<blitz::Array<T,1> >
{
protected:
  typedef blitz::Array<T,1> T_lambda;

  /// The length of the wavelength vector.
  int n;

public:
  isotropic_angular_distribution () {};
  isotropic_angular_distribution (const T_lambda& l):
    angular_distribution<T_lambda> (), n(l.size()) {};

  // assignment operator implicitly generated

  /** Returns the distribution function for an isotropic distribution,
      which is just 1/4pi.  To save space, we do not keep an array of
      the constant number, we create the array on-the-fly each time.
      Both alternatives are bad, but it is unfortunately not possible
      to return a constant since we are overloading a virtual
      function.  */
  virtual T_lambda distribution_function (const vec3d&, const vec3d&) const {
    T_lambda d(n);
    isotropic_angular_distribution::distribution_function(vec3d(), vec3d(), d);
    return d;};
  /** Writes the distribution function to the supplied array, avoiding
      allocation of a temporary. */
  virtual void distribution_function (const vec3d&, const vec3d&,
				      T_lambda& distr) const {
    distr =  1./(4*constants::pi); };
  boost::shared_ptr<angular_distribution<blitz::Array<T,1> > > clone() const {
    return boost::shared_ptr<angular_distribution<blitz::Array<T,1> > >
      (new isotropic_angular_distribution<blitz::Array<T,1> >(*this));};
};


/** Angular_distribution class for collimated rays.  Simply returns 1
    for the set direction, otherwise 0. (This is not really valid, as
    it should be a delta function, but in any case the distribution is
    nonzero only on a set of measure zero, so it can't really be
    sampled with MC techniques.) */
template <typename T_lambda> 
class mcrx::collimated_angular_distribution:
  public angular_distribution<T_lambda>
{
  const vec3d direction_;

public:
  collimated_angular_distribution (const vec3d& dir, const T_lambda& l):
    angular_distribution<T_lambda> (), direction_ (dir) {};

  const vec3d& direction () const {return direction_;};
  /// The distribution function.
  virtual T_lambda distribution_function (const vec3d& dir, const vec3d&) const {
    // The first direction vector is the important one, since the
    // second one is only used when scattering
    return (dot (dir, direction_) - 1) >-1e-9?  1:0;};
  virtual void distribution_function (const vec3d& dir, const vec3d&,
				      T_lambda& distr) const {
    distr = (dot (dir, direction_) - 1) >-1e-9?  1:0; };
  boost::shared_ptr<angular_distribution<T_lambda> > clone() const {
    return boost::shared_ptr<angular_distribution<T_lambda> >
      (new collimated_angular_distribution<T_lambda>(*this));};
};

/** Specialization of the collimated_angular_distribution class for
    polychromatic situations. */
template <typename T> 
class mcrx::collimated_angular_distribution<blitz::Array<T,1> >: 
  public angular_distribution<blitz::Array<T,1> >
{
  typedef blitz::Array<T,1> T_lambda;

  /// The length of the wavelength vector.
  const int n;
  const vec3d direction_;

public:
  collimated_angular_distribution (const vec3d& dir, const T_lambda& l):
    angular_distribution<T_lambda> (), direction_ (dir), n (l.size())  {};

  const vec3d& direction () const {return direction_;};
  /// The distribution function.
  virtual T_lambda distribution_function (const vec3d& dir, const vec3d&) const {
    // The first direction vector is the important one, since the
    // second one is only used when scattering
    T_lambda d(n); 
    collimated_angular_distribution::distribution_function(dir, vec3d(), d);
    return d;};
  virtual void distribution_function (const vec3d& dir, const vec3d&,
				      T_lambda& distr) const {
    distr = (dot (dir, direction_) - 1) >-1e-9?  1:0; };

  boost::shared_ptr<angular_distribution<blitz::Array<T,1> > > clone() const {
    return boost::shared_ptr<angular_distribution<blitz::Array<T,1> > >
      (new collimated_angular_distribution<blitz::Array<T,1> >(*this));};
};


/** Angular_distribution class for an isotropic distribution within a
    certain angle around a specified direction. */
template <typename T_lambda> 
class mcrx::cone_angular_distribution: 
  public angular_distribution<T_lambda>
{
private:
  const vec3d direction_;
  T_float cos_theta_; ///< Cos of the opening angle.
  T_float value_; ///< The distribution value is 1./total solid angle subtended.

public:
  cone_angular_distribution (const vec3d& dir, T_float theta,
			     const T_lambda& l): 
    angular_distribution<T_lambda>(), 
    direction_(dir), 
    cos_theta_(cos(theta)), 
    value_(1./(4*constants::pi*sin(theta/2)*sin(theta/2))) {};
  const vec3d& direction () const {return direction_;};
  /** Returns the distribution function for a cone angular
      distribution, which is 1/omega, where omega is the total solid
      angle subtended by the cone, if we are within the cone and zero
      otherwise. */
  virtual T_lambda distribution_function (const vec3d& dir, const vec3d&) const {
    // The first direction vector is the important one, since the
    // second one is only used when scattering
    const T_float costheta = dot(dir,direction_);
    return costheta > cos_theta_ ? value_ : 0;
  };
  virtual void distribution_function (const vec3d& dir, const vec3d&,
				      T_lambda& distr) const {
    const T_float costheta = dot(dir,direction_);
    distr = (costheta > cos_theta_) ? value_ : 0;
  };

  boost::shared_ptr<angular_distribution<T_lambda> > clone() const {
    return boost::shared_ptr<angular_distribution<T_lambda> >
      (new cone_angular_distribution<T_lambda>(*this));};
};


/** Specialization of the cone_angular_distribution class for
    polychromatic situations. */
template <typename T> 
class mcrx::cone_angular_distribution<blitz::Array<T,1> >: 
  public angular_distribution<blitz::Array<T,1> >
{
protected:
  typedef blitz::Array<T,1> T_lambda;
  const vec3d direction_;
  T_float cos_theta_; ///< Cos of the opening angle.
  T_float value_; ///< The distribution value is 1./total solid angle subtended.

  /// The length of the wavelength vector.
  const int n;

public:
  cone_angular_distribution (const vec3d& dir, T_float theta,
			     const T_lambda& l): 
    angular_distribution<T_lambda> (), 
    direction_(dir), 
    cos_theta_(cos(theta)), 
    value_(1./(4*constants::pi*sin(theta/2)*sin(theta/2))),
    n(l.size()) {};

  const vec3d& direction () const {return direction_;};
  /** Returns the distribution function for a cone angular
      distribution, which is 1/omega, where omega is the total solid
      angle subtended by the cone, if we are within the cone and zero
      otherwise. */
  virtual T_lambda distribution_function (const vec3d& dir, const vec3d&) const {
    T_lambda d(n); 
    cone_angular_distribution::distribution_function(dir, vec3d(), d);
    return d;
  };
  virtual void distribution_function (const vec3d& dir, const vec3d&,
				      T_lambda& distr) const {
    // The first direction vector is the important one, since the
    // second one is only used when scattering
    const T_float costheta = dot(dir,direction_);
    distr = (costheta > cos_theta_) ? value_ : 0;
  };

  boost::shared_ptr<angular_distribution<blitz::Array<T,1> > > clone() const {
    return boost::shared_ptr<angular_distribution<blitz::Array<T,1> > >
      (new cone_angular_distribution<blitz::Array<T,1> >(*this));};
};


/** A composition of a ray, a normalization vector, an
    angular_distribution object, and a velocity vector. (Name is for
    historical reasons, the normalization vector was added later...)
    Is used as a return type when scattering or emitting rays, as all
    this information must be returned. Note that the ray is returned
    with an uninitialized column density array, because the emission
    classes don't know anything about column densities. */
template <typename T_lambda>
class mcrx::ray_distribution_pair {
private:
  boost::shared_ptr<ray<T_lambda> > rayptr_;
  T_lambda norm_;

  // can't keep a reference here because sometimes the distribution is
  // a temporary
  boost::shared_ptr<angular_distribution<T_lambda> > distr_;
  vec3d vel_;

public:
  ray_distribution_pair (const boost::shared_ptr<ray<T_lambda> >& f,
			 const T_lambda& n,
                         const angular_distribution<T_lambda>& s,
			 const vec3d& v=vec3d(0,0,0)):
    rayptr_ (f), norm_(make_thread_local_copy(n)), distr_ (s.clone()), vel_(v) {};

  ray_distribution_pair (const boost::shared_ptr<ray<T_lambda> >& f,
			 const T_lambda& n,
                         boost::shared_ptr<angular_distribution<T_lambda> >& s,
			 const vec3d& v=vec3d(0,0,0)):
    rayptr_ (f), norm_(make_thread_local_copy(n)), distr_ (s), vel_(v)  {};

  ray_distribution_pair(const ray_distribution_pair& rhs) :
    rayptr_(rhs.rayptr_), norm_(rhs.norm_), 
    distr_(rhs.distr_), vel_(rhs.vel_) {};

  boost::shared_ptr<ray<T_lambda> >& ray_pointer() {return rayptr_;};
  T_lambda& normalization() {return norm_;};

  const boost::shared_ptr<angular_distribution<T_lambda> >& 
    distribution() const {return distr_;};

  const vec3d& velocity() const {return vel_;};
};
		      
  
#endif
